import os
import struct
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 读取数据
def load_mnist(path, kind='train'):
    """Load MNIST data from `path`"""
    labels_path = path + '/' + kind + '-labels.idx1-ubyte'
    images_path = path + '/' + kind + '-images.idx3-ubyte'

    with open(labels_path, 'rb') as lbpath:
        magic, n = struct.unpack('>II',
                                 lbpath.read(8))
        labels = np.fromfile(lbpath,
                             dtype=np.uint8)

    with open(images_path, 'rb') as imgpath:
        magic, num, rows, cols = struct.unpack('>IIII',
                                               imgpath.read(16))
        images = np.fromfile(imgpath,
                             dtype=np.uint8).reshape(len(labels), 784)

    return images, labels

def visual_mnist(images, labels):
    # 查看mnist数据集的图片
    fig, ax = plt.subplots(
        nrows=2,
        ncols=5,
        sharex=True,
        sharey=True, )

    ax = ax.flatten()
    for i in range(10):
        img = images[labels == i][0].reshape(28, 28)
        ax[i].imshow(img, cmap='Greys', interpolation='nearest')

    ax[0].set_xticks([])
    ax[0].set_yticks([])
    plt.tight_layout()
    plt.show()


# 模型定义
class Perceptron:

    # 原始形式算法
    # w_dims 表示特征个数
    def __init__(self,w_dims,rate=0.1): # 初始化

        self.w = np.zeros(w_dims)
        self.b = 0
        self.rate = rate

        self.bestScore = 0
        self.bestw = np.zeros(w_dims)
        self.bestb = 0

    def sign(self,x): # 前向传播

        pred = np.sign( np.matmul(self.bestw, x) + self.bestb )

        return pred

    def fit(self, x, y, x_test, y_test): # 训练
        
        AllPass = False
        while_counts = 0
        while AllPass is not True:
            
            AllPass = True
            for i in range(x.shape[0]):

                flag = y[i] * ( np.matmul(self.w, x[i]) + self.b )
                if flag <= 0:
                    self.w = self.w + self.rate * y[i] * x[i]
                    self.b = self.b + self.rate * y[i]
                    AllPass = False

            while_counts += 1
            if while_counts % 10 == 0:
                score_train = self.score(x, y)
                score_test  = self.score(x_test, y_test)
                print('The data set has been traversed {:5d} times, Training score: {:7.2%}, Testing score: {:7.2%}'
                        .format(while_counts, score_train, score_test))
                if score_test > self.bestScore:
                    self.bestScore = score_test
                    self.bestw = self.w
                    self.bestb = self.b

            if while_counts > 10000:
                AllPass = True

        print('Mission Complete!')

    def score(self,x,y): # 测试

        total_num = 0
        correct_num = 0
        for i in range(x.shape[0]):
            flag = y[i] * ( np.matmul(self.w, x[i]) + self.b )
            if flag > 0:
                correct_num += 1
            total_num += 1

        return correct_num/total_num

class Perceptron_Duality:

    # 对偶形式算法
    # x_nums表示训练集中元素的个数
    def __init__(self,x_nums,rate=0.1): # 初始化

        self.a = np.zeros(x_nums)
        self.b = 0
        self.rate = rate
        # w 要在训练结束后更新，不然无法在没有y的情况下前向传播
        self.w = 0

        self.bestScore = 0
        self.bestw = 0
        self.bestb = 0

    def sign(self,x): # 前向传播

        pred = np.sign( np.matmul(self.w, x) + self.b )

        return pred

    def fit(self, x, y, x_test, y_test): # 训练
        
        # 构建Gram内积矩阵
        # Gram = np.matmul( x, np.transpose(x) )
        # np.save('Gram_mnist.npy',Gram)
        Gram = np.load('Gram_mnist.npy',allow_pickle=True)
        print('Gram Complete！')

        AllPass = False
        while_counts = 0
        while AllPass is not True:

            AllPass = True
            for i in range(x.shape[0]):

                flag = y[i] * ( np.dot( self.a, y*Gram[:,i] ) + self.b )
                if flag <= 0:
                    self.a[i] = self.a[i] + self.rate 
                    self.b = self.b + self.rate * y[i]
                    AllPass = False

            self.w = np.matmul( self.a*y, x )

            while_counts += 1
            if while_counts % 10 == 0:
                score_train = self.score(x, y)
                score_test  = self.score(x_test, y_test)
                print('The data set has been traversed {:5d} times, Training score: {:7.2%}, Testing score: {:7.2%}'
                        .format(while_counts, score_train, score_test))
                if score_test > self.bestScore:
                    self.bestScore = score_test
                    self.bestw = self.w
                    self.bestb = self.b

            if while_counts > 10000:
                AllPass = True
        
        print('Mission Complete!')

    def score(self,x,y): # 测试
        
        total_num = 0
        correct_num = 0
        for i in range(x.shape[0]):
            flag = y[i] * ( np.matmul(self.w, x[i]) + self.b )
            if flag > 0:
                correct_num += 1
            total_num += 1

        return correct_num/total_num

# 主函数
def main():

    # 加载数据，打印基本信息
    images_train, labels_train = load_mnist('../Datasets/mnist')
    print('Images_train Shape:{:<25s} Labels_train Shape:{:<25s}'
            .format(str(images_train.shape),str(labels_train.shape)))
    print('Images_train type :{:<25s} Labels_train type :{:<25s}'
            .format(str(type(images_train)),str(type(labels_train))))
    images_test, labels_test = load_mnist('../Datasets/mnist','t10k')
    print('Images_test Shape :{:<25s}  Labels_test Shape:{:<25s}'
            .format(str(images_test.shape),str(labels_test.shape)))
    print('Images_test type  :{:<25s}  Labels_test type :{:<25s}'
            .format(str(type(images_test)),str(type(labels_test))))
    
    # 处理数据
    labels_train_new = np.array([1 if y == 1 else -1 for y in labels_train])
    labels_test_new  = np.array([1 if y == 1 else -1 for y in labels_test ])

    # 创建模型与训练
    model = Perceptron(images_train.shape[1])
    # model = Perceptron_Duality(images_train.shape[0]) # 对偶形式要比原始形式的慢，可能是因为矩阵太大了索引很慢
    model.fit(images_train, labels_train_new, images_test, labels_test_new)
    print(model.bestw.shape,model.bestb,model.bestScore)

if __name__ == '__main__':
    main()